import argparse

def convert_dataset(prefix, input_output_path):
    # 读取容量
    with open(f"../originalData/{prefix}/{prefix}_c.txt") as f:
        c = int(f.read().strip())
    
    # 读取重量和利润
    with open(f"../originalData/{prefix}/{prefix}_w.txt") as f:
        weights = [int(line.strip()) for line in f]
    
    with open(f"../originalData/{prefix}/{prefix}_p.txt") as f:
        profits = [int(line.strip()) for line in f]
    
    # 读取选中状态
    with open(f"../originalData/{prefix}/{prefix}_s.txt") as f:
        selected = [int(line.strip()) for line in f]
    
    # 验证数据一致性
    n = len(weights)
    if len(profits) != n or len(selected) != n:
        raise ValueError("Files have inconsistent number of items")
    
    # 生成input.txt内容
    input_lines = [f"{n} {c}"]
    for p, w in zip(profits, weights):
        input_lines.append(f"{p} {w}")
    
    # 生成output.txt内容
    total_profit = sum(p for p, s in zip(profits, selected) if s == 1)
    selected_ids = [i+1 for i, s in enumerate(selected) if s == 1]
    
    # 写入文件
    with open(f"{input_output_path}/input.txt", "w") as f:
        f.write("\n".join(input_lines))
    
    with open(f"{input_output_path}/output.txt", "w") as f:
        f.write(f"{total_profit}\n")
        f.write(" ".join(map(str, selected_ids)))

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Convert knapsack dataset')
    parser.add_argument('--prefix', type=str, required=True, help='File prefix (e.g. p01)')
    parser.add_argument('--output', type=str, required=True, help='Output directory path')
    
    args = parser.parse_args()
    
    try:
        convert_dataset(args.prefix, args.output)
        print(f"Conversion successful! Files saved to {args.output}")
    except Exception as e:
        print(f"Error: {str(e)}")
        exit(1)